from SQLiteDatabase import SQLiteDatabase
import json

create_table_sqls = {
    "sample": """
        CREATE TABLE IF NOT EXISTS sample (
        Id INTEGER PRIMARY KEY AUTOINCREMENT,    -- 主键，自增
        operator_id INTEGER NOT NULL,            -- 样本上传者ID
        operator_name VARCHAR(32) NOT NULL,      -- 样本上传者
        upload_time DATETIME NOT NULL,           -- 样本上传时间
        is_delete TINYINT NOT NULL,              -- 是否逻辑删除
        origin_name VARCHAR(128) NOT NULL,       -- 样本原始名称
        md5 VARCHAR(128) NOT NULL                -- 样本md5值
        )
        """,
    "sample_info": """
            CREATE TABLE IF NOT EXISTS sample_info (
            Id INTEGER PRIMARY KEY AUTOINCREMENT,     -- 主键，自增
            md5 VARCHAR(128) NOT NULL,                -- 样本md5值
            fam_label VARCHAR(128),                   -- 家族标签信息
            graph_embedding  VARCHAR(128),            -- 恶意软件的嵌入向量
            is_train_sample TINYINT NOT NULL          -- 是否为训练用的样本
        )""",
    "sim_score": """
            CREATE TABLE IF NOT EXISTS sim_score (
            Id INTEGER PRIMARY KEY AUTOINCREMENT,     -- 主键，自增
            base_sample INTEGER NOT NULL,             -- 基准样本ID
            target_samples VARCHAR(128) NOT NULL,     -- 待对比的样本列表
            operator_name VARCHAR(128) NOT NULL,      -- 操作人姓名
            modify_time DATETIME NOT NULL,            -- 操作时间
            sim_scores VARCHAR(128) NOT NULL,         -- 相似性度量结果
            sim_type TINYINT NOT NULL                 -- 相似性度量类型
        )""",
}

if __name__ == "__main__":
    with open("config.json", "r") as file:
        config = json.load(file)
    db = SQLiteDatabase(config["database"]["dbName"])
    for sql in create_table_sqls:
        db.create_table(create_table_sqls[sql])
